package org.znxs.znagent_s.chatMemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的对话记忆
 */
public class FileBasedChatMemory implements ChatMemory {
    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        // 设置是否必须实例化
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());

    }

    // 构建对象的时候指定文件保存目录
    public FileBasedChatMemory(String baseDir) {
        BASE_DIR = baseDir;
        File file = new File(BASE_DIR);
        if (!file.exists()) {
            boolean mkdirs = file.mkdirs();
            if (!mkdirs) {
                throw new RuntimeException("创建文件夹失败: " + BASE_DIR);
            }
        }
    }


    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        /// 从allMessageList列表中提取最近的 lastN 条消息
        List<Message> allMessageList = getOrCreateConversation(conversationId);
        return allMessageList
                .stream()
                .skip(Math.max(0, allMessageList.size() - lastN)) // 用于跳过前N个元素
                .toList();
    }


    @Override
    public void clear(String conversationId) {
        // 直接销毁文件就行了
        File conversationFile = getConversationFile(conversationId);
        if (conversationFile.exists()) {
            conversationFile.delete();
        }
    }


    /**
     * 从文件中获取集合
     *
     * @param conversationId 对话id
     * @return
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        // 根据对话id获取指定对话的存储文件
        File conversationFile = getConversationFile(conversationId);
        List<Message> messageList = new ArrayList<>();
        if (conversationFile.exists()) {
            try (Input input = new Input(new FileInputStream(conversationFile))) {
                // 创建输入流，将文件内容读取到集合中
                messageList = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        // 保存之后 进行返回
        return messageList;

    }

    private void saveConversation(String conversationId, List<Message> messageList) {
        // 获取文件
        File conversationFile = getConversationFile(conversationId);
        // 文件一定是存在的，所以直接写入
        try (Output output = new Output(new FileOutputStream(conversationFile))) {
            // 创建输出流，将集合写入到存储文件中
            kryo.writeObject(output, messageList);
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    /**
     * 获取对应的对话记忆存储集合 保存文件
     *
     * @param conversationId 对话id
     * @return
     */
    private File getConversationFile(String conversationId) {
        // 采用kryo的高性能存储
        return new File(BASE_DIR, conversationId + ".kryo");
    }


}
